from data import ResearcherState
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command
from typing import Literal
from utils import get_all_tools
import asyncio
from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
    filter_messages,
    get_buffer_string,
)
from utils import get_today_str

# Tool Execution Helper Function
async def execute_tool_safely(tool, args, config):
    """Safely execute a tool with error handling."""
    try:
        return await tool.ainvoke(args, config)
    except Exception as e:
        return f"Error executing tool: {str(e)}"


async def researcher_tools(state: ResearcherState, config: RunnableConfig) -> Command[Literal["researcher", "compress_research"]]:
    """Execute tools called by the researcher, including search tools and strategic thinking.
    
    This function handles various types of researcher tool calls:
    1. think_tool - Strategic reflection that continues the research conversation
    2. Search tools (tavily_search, web_search) - Information gathering
    3. MCP tools - External tool integrations
    4. ResearchComplete - Signals completion of individual research task
    
    Args:
        state: Current researcher state with messages and iteration count
        config: Runtime configuration with research limits and tool settings
        
    Returns:
        Command to either continue research loop or proceed to compression
    """
    # Step 1: Extract current state and check early exit conditions
    # configurable = Configuration.from_runnable_config(config)
    researcher_messages = state.get("researcher_messages", [])
    most_recent_message = researcher_messages[-1]
    
    # Early exit if no tool calls were made (including native web search)
    has_tool_calls = bool(most_recent_message.tool_calls)
    # has_native_search = (
    #     openai_websearch_called(most_recent_message) or 
    #     anthropic_websearch_called(most_recent_message)
    # )
    
    # if not has_tool_calls and not has_native_search:
    if not has_tool_calls:
        print(f"++++++++++++++++++ researcher_tools: goto compress_research without tools calls ++++++++++++++++++++++")
        return Command(goto="compress_research")
    
    # Step 2: Handle other tool calls (search, MCP tools, etc.)
    tools = await get_all_tools(config)
    tools_by_name = {
        tool.name if hasattr(tool, "name") else tool.get("name", "web_search"): tool 
        for tool in tools
    }
    
    # Execute all tool calls in parallel
    tool_calls = most_recent_message.tool_calls
    tool_execution_tasks = [
        execute_tool_safely(tools_by_name[tool_call["name"]], tool_call["args"], config) 
        for tool_call in tool_calls
    ]
    print(f"++++++++++++++++++ researcher_tools: execute tools ++++++++++++++++++++++")
    print(f"tool_calls: {tool_calls}")

    observations = await asyncio.gather(*tool_execution_tasks)
    
    # Create tool messages from execution results
    tool_outputs = [
        ToolMessage(
            content=observation,
            name=tool_call["name"],
            tool_call_id=tool_call["id"]
        ) 
        for observation, tool_call in zip(observations, tool_calls)
    ]
    
    # Step 3: Check late exit conditions (after processing tools)
    # exceeded_iterations = state.get("tool_call_iterations", 0) >= configurable.max_react_tool_calls
    exceeded_iterations = state.get("tool_call_iterations", 0) >= 2
    research_complete_called = any(
        tool_call["name"] == "ResearchComplete" 
        for tool_call in most_recent_message.tool_calls
    )
    
    if exceeded_iterations or research_complete_called:
        # End research and proceed to compression
        print(f"++++++++++++++++++ researcher_tools: goto compress_research ++++++++++++++++++++++")
        print(f"goto compress_research because of exceeded_iterations: {exceeded_iterations} or research_complete_called: {research_complete_called} ")
        print(f"tool_outputs: {tool_outputs}")
        return Command(
            goto="compress_research",
            update={"researcher_messages": tool_outputs}
        )
    
    # Continue research loop with tool results
    print(f"++++++++++++++++++ researcher_tools: goto researcher with tool_outputs ++++++++++++++++++++++")
    # print(tool_outputs)
    return Command(
        goto="researcher",
        update={"researcher_messages": tool_outputs}
    )

